# 加载环境变量
from dotenv import load_dotenv, find_dotenv
_ = load_dotenv(find_dotenv())

# 自研包导入
from AutoAgent.AutoAgent import AutoAgent
from Tools.Tools import *
from Tools.DatabaseQueryTool import load_data, DatabaseQuerist

# 三方包导入
from langchain_openai import ChatOpenAI
# from langchain_community.vectorstores import Chroma
from langchain_chroma import Chroma
# from langchain_community.document_loaders import DirectoryLoader, JSONLoader
from langchain_huggingface import HuggingFaceEmbeddings
import gradio as gr
from gradio import ChatMessage

import os
from sys import argv
from pathlib import Path
import time

STATUS_OK = 0
STATUS_ERR = -1


def launch_agent(agent: AutoAgent, 
                 web_ui, 
                 examples_path="./prompts/main", 
                 examples_file="examples.txt"):
    """执行 agent 函数"""
    def chat_func(message, history):
        start_time = time.time()
        output = []

        phase = "thinking process"
        content = ""
        output.append(ChatMessage(
            content=content,
            metadata={"title": "思考过程", "id": 0, "status": "pending"}
        ))
        yield output

        flag = False
        for s in agent.run(message, verbose=True, web_ui=True):
            if "\n----\nFINISH\n" == s:
                content += s
                output[0].content = content
                flag = True
                phase = "reply"
                content = ""
                output.append(ChatMessage(content=content))
            elif "抱歉!我没能完成您的任务。" == s:
                flag = True
                phase = "reply"
                content = s
                output.append(ChatMessage(content=content))
            else:
                content += s

            if flag:
                output[0].metadata["status"] = "done"
                output[0].metadata["duration"] = time.time() - start_time
                flag = False
            else:
                if "thinking process" == phase:
                    output[0].content = content
                else:
                    output[1].content = content

            yield output

    if web_ui:
        examples = Path(os.path.join(examples_path, examples_file)).read_text().split("\n")

        gr.ChatInterface(
            chat_func,
            type="messages",
            textbox=gr.Textbox(placeholder="输入查询...", submit_btn=True, stop_btn=True),
            title="酒店信息智能查询系统",
            description="本系统支持根据名称、类型、地址、地铁站、电话号码、设施、价格区间、评分等查询酒店信息，并支持多轮对话。",
            theme="default",
            examples=examples,
            cache_examples=False,
            save_history=True,
        ).queue().launch()
    else:
        human_icon = "\U0001F468"
        ai_icon = "\U0001F916"

        # 对话程序的主循环
        while True:
            task = input(f"{ai_icon}：有什么可以帮您？\n{human_icon}：")
            # 退出条件
            if task.strip().lower() == "quit":
                break
            # reply = agent.run(task, verbose=True)
            # print(f"{ai_icon}：{reply}\n")
            print(f"{ai_icon}：", end="")
            for s in agent.run(task, verbose=True):
                print(s, end="")
            print("\n", end="\n")

def main(argv):
    """主函数"""
    argc = len(argv)

    if argc >= 2:
        web_ui = argv[1]
    else:
        web_ui = "true"

    if "true" == web_ui.lower():
        web_ui = True
    elif "false" == web_ui.lower():
        web_ui = False
    else:
        print(f"[Err] The boolean value (\"{web_ui}\") is invalid!")
        return STATUS_ERR

    # 加载数据
    load_data()

    # 语言模型
    embedding_function = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")

    # 存储长时记忆的向量数据库
    db = Chroma(embedding_function=embedding_function)
    retriever = db.as_retriever(
        search_type="similarity_score_threshold",
        search_kwargs={"k": 1, "score_threshold": 0.1}
    )

    # 自定义工具集
    tools = [
        # database_query_tool,
        DatabaseQuerist(prompts_path="./prompts/tools",
                        prompt_file_1="database_querist_1.json",
                        prompt_file_2="database_querist_2.json",
                        verbose=True,
                        web_ui=web_ui).as_tool(),
        finish_placeholder
    ]

    # 语言模型，agent 大脑
    llm = ChatOpenAI(
        # model="gpt-4-1106-preview",
        model="deepseek-chat",
        temperature=0,
        seed=42
    )

    # 定义智能体
    agent = AutoAgent(
        # 大脑
        llm=llm,
        # 提示词的目录
        prompts_path="./prompts/main",
        # 可用的工具
        tools=tools,
        # 工作目录
        work_dir="./data",
        # 主要提示词
        main_prompt_file="main.json",
        # 任务结束的提示词
        final_prompt_file="final_step.json",
        # 最大的思考轮数
        max_thought_steps=20,
        # 长期记忆
        memery_retriever=retriever
    )

    # 运行智能体
    launch_agent(agent, web_ui)

    return STATUS_OK


if __name__ == "__main__":
    main(argv)
